from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START
from langgraph.checkpoint.memory import InMemorySaver

# 构建deepseek大模型客户端
llm = ChatOpenAI(
    model="deepseek-chat",
    base_url="https://api.deepseek.com/v1",
    openai_api_key="sk-ddad02a64c4c47ea9f6f526ef47cb602",
)


def call_model(state: MessagesState):
    """
    调用大语言模型处理消息状态

    Args:
        state (MessagesState): 包含对话消息的状态对象，其中messages字段存储对话历史

    Returns:
        dict: 返回包含模型响应消息的字典，格式为{"messages": response}
    """
    response = llm.invoke(state["messages"])
    return {"messages": response}


# 创建状态图构建器，使用MessagesState作为状态类型
builder = StateGraph(MessagesState)

# 添加节点和边来构建工作流图
builder.add_node(call_model)
builder.add_edge(START, "call_model")

# 创建内存检查点保存器，用于保存对话状态
checkpointer = InMemorySaver()

# 编译图并设置检查点保存器
graph = builder.compile(checkpointer=checkpointer)
# graph = builder.compile()

# 配置字典，用于存储应用程序的配置信息
# configurable: 包含可配置项的字典
#   thread_id: 线程ID标识符，用于区分不同的执行线程
config = {
    "configurable": {
        "thread_id": "1"
    }
}

# 第一轮对话：询问湖南的省会
# 遍历图模型的流式响应chunks，并打印每条消息
# 该代码块通过graph.stream方法向模型发送用户查询消息，
# 并以流式方式获取响应结果。每次迭代处理一个响应chunk，
# 最终打印出最后一条消息的内容。
# 参数说明:
#   graph: 图模型对象，需具备stream方法
#   config: 配置信息，用于控制模型行为
#   stream_mode: 流式传输模式，设置为"values"表示按值传输
# 消息格式:
#   发送的消息包含角色(user)和内容(湖南的省会是哪⾥？)
#   响应的消息通过chunk["messages"]获取，其中最后一条消息为最新回复

for chunk in graph.stream(
        {"messages": [{"role": "user", "content": "湖南的省会是哪⾥？"}]},
        config,
        stream_mode="values",
):
    # 打印当前chunk中最后一条消息的格式化内容
    chunk["messages"][-1].pretty_print()

# 第二轮对话：询问湖北的省会（基于之前的对话上下文）
for chunk in graph.stream(
        {"messages": [{"role": "user", "content": "湖北呢？"}]},
        config,
        stream_mode="values",
):
    chunk["messages"][-1].pretty_print()
